package unionFind

/*
参考优质文章：https://zhuanlan.zhihu.com/p/93647900/
*/

type UnionFind2 struct {
	// 路径压缩，现在u.Parent[x]不再表示x的父节点，而是表示x所在集合的代表元素
	Parent []int
	// 按秩合并，初始化时将所有节点的秩设置为1（可以理解为树的深度,根节点深度为1）
	Rank []int
}

func (u *UnionFind2) Init(n int) {
	u.Parent = make([]int, n)
	u.Rank = make([]int, n)
	for i := 0; i < n; i++ {
		u.Parent[i] = i
		u.Rank[i] = 1
	}
}

// 查找元素x所在集合的代表元素
// 路径压缩，现在u.Parent[x]不再表示x的父节点，而是表示x所在集合的代表元素
func (u *UnionFind2) Find(x int) int {
	if x == u.Parent[x] {
		return x
	}
	u.Parent[x] = u.Find(u.Parent[x])
	return u.Parent[x]
}

// 判断元素i,j是否在同一个集合中
func (u *UnionFind2) IsConnected(i, j int) bool {
	return u.Find(i) == u.Find(j)
}

// 将i所在的集合合并到j所在的集合
func (u *UnionFind2) Merge(i, j int) {
	rootX, rootY := u.Find(i), u.Find(j)
	// 当两颗树秩相同时，将i所在的树合并到j所在的树
	// 当一棵树的秩小于另一棵树时，将秩小的树合并到秩大的树，因为如果将秩大的合并到秩小的，则合并后树的深度会加深，使得每个节点到根节点距离变长
	// 如果将秩小的树合并到秩大的树，则树的深度不变；因次将秩小的树合并到秩大的树
	if u.Rank[rootX] <= u.Rank[rootY] {
		u.Parent[rootX] = rootY
	} else {
		u.Parent[rootY] = rootX
	}
	// 当两颗树的秩相同且i与j不在同一个集合中时，合并后秩将+1
	if u.Rank[rootX] == u.Rank[rootY] && rootX != rootY {
		u.Rank[rootY]++
	}
}
